package mr

import (
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net/rpc"
	"os"
	"sort"
)

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

// Map functions return a slice of KeyValue.
type KeyValue struct {
	Key   string
	Value string
}

// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

// main/mrworker.go calls this function.
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.

	// uncomment to send the Example RPC to the coordinator.
	// CallExample()

	for {
		var task TaskAssignment
		ok := call("Coordinator.AssignTask", TaskResult{}, &task)
		if !ok || len(task.DataPath) <= 0 {
			continue
		}

		switch {
		case task.Type == Map:
			intermediateKV, err := doMap(mapf, task.DataPath[0], task.NReduce)
			if err != nil {
				log.Printf("map task %d failed to finish map: %v", task.Id, err)
				workerFailed(&task)
				// break
				continue
			}
			filenames, err := sinkMap(intermediateKV, task.Id)
			if err != nil {
				log.Printf("map task %d failed to sink map: %v", task.Id, err)
				workerFailed(&task)
				// break
				continue
			}
			call("Coordinator.ReturnTaskResult", TaskResult{Id: task.Id, Type: task.Type, ResultPath: filenames, DataPath: []string{}}, &TaskAssignment{})

		case task.Type == Reduce:
			resultKV, err := doReduce(reducef, task.DataPath)
			if err != nil {
				log.Printf("reduce task %d failed to finish reduce: %v", task.Id, err)
				workerFailed(&task)
				// break
				continue
			}
			filename, err := sinkReduce(resultKV, task.Id)
			if err != nil {
				log.Printf("reduce task %d failed to sink reduce: %v", task.Id, err)
				workerFailed(&task)
				// break
				continue
			}
			call("Coordinator.ReturnTaskResult", TaskResult{Id: task.Id, Type: task.Type, ResultPath: []string{filename}, DataPath: []string{}}, &TaskAssignment{})

		case task.Type == Exit:
			return
		}

	}

}

// send coordinator the failed into via RPC
func workerFailed(task *TaskAssignment) {
	call("Coordinator.ReturnTaskResult", task.ToTaskResult(), &TaskAssignment{})
}

// do map.
// input the map function, the file to read and split the result to R parts.
// return the map result(split to R parts) and error
func doMap(mapf func(string, string) []KeyValue, filename string, R int) ([][]KeyValue, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	content, err := io.ReadAll(file)
	if err != nil {
		return nil, err
	}
	file.Close()
	kva := mapf(filename, string(content))

	res := make([][]KeyValue, R)
	for _, kv := range kva {
		idx := ihash(kv.Key) % R
		res[idx] = append(res[idx], kv)
	}
	return res, nil
}

// Sink reduce result. Write to tmp files first and rename them
func sinkMap(intermediateKV [][]KeyValue, taskId int) ([]string, error) {
	resFilename := make([]string, len(intermediateKV))
	tmpFilename := make([]string, len(intermediateKV))

	for idx, kvs := range intermediateKV {
		if kvs == nil {
			continue
		}
		filename := fmt.Sprintf("mr-%d-%d", taskId, idx)
		tmpname := fmt.Sprintf("tmp-%s", filename)

		intermediateFile, err := os.Create(tmpname)
		if err != nil {
			return nil, err
		}
		jsonEncoder := json.NewEncoder(intermediateFile)
		for _, kv := range kvs {
			err := jsonEncoder.Encode(&kv)
			if err != nil {
				return nil, err
			}
		}

		resFilename[idx] = filename
		tmpFilename[idx] = tmpname
	}

	for idx := range intermediateKV {
		if tmpFilename[idx] == "" {
			continue
		}
		rename_err := os.Rename(tmpFilename[idx], resFilename[idx])
		if rename_err != nil {
			return nil, rename_err
		}
	}

	return resFilename, nil
}

// do map.
// input the reduce function and files to read.
// return the reduce result and error
func doReduce(reducef func(string, []string) string, files []string) ([]KeyValue, error) {
	var kvs, res []KeyValue
	for _, filename := range files {
		if filename == "" {
			continue
		}
		fd, err := os.Open(filename)
		if err != nil {
			return nil, err
		}
		jsonDecoder := json.NewDecoder(fd)
		for {
			var kv KeyValue
			if err := jsonDecoder.Decode(&kv); err != nil {
				break
			}
			kvs = append(kvs, kv)
		}
	}
	if kvs == nil {
		return nil, nil
	}

	sort.Sort(ByKey(kvs))

	i := 0
	kvsLength := len(kvs)
	for i < kvsLength {
		j := i + 1
		for j < kvsLength && kvs[j].Key == kvs[i].Key {
			j++
		}
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, kvs[k].Value)
		}
		output := reducef(kvs[i].Key, values)
		res = append(res, KeyValue{kvs[i].Key, output})
		i = j
	}

	return res, nil
}

// Sink reduce result. Write to the tmp file first and rename it
func sinkReduce(resultKV []KeyValue, taskId int) (string, error) {
	if resultKV == nil {
		return "", nil
	}

	filename := fmt.Sprintf("mr-out-%d", taskId)
	tmpname := fmt.Sprintf("tmp-%s", filename)
	fd, err := os.Create(tmpname)
	if err != nil {
		return "", err
	}
	for _, kv := range resultKV {
		fmt.Fprintf(fd, "%v %v\n", kv.Key, kv.Value)
	}
	fd.Close()
	rename_err := os.Rename(tmpname, filename)
	if rename_err != nil {
		return "", rename_err
	}

	return filename, nil
}

// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	// the "Coordinator.Example" tells the
	// receiving server that we'd like to call
	// the Example() method of struct Coordinator.
	ok := call("Coordinator.Example", &args, &reply)
	if ok {
		// reply.Y should be 100.
		fmt.Printf("reply.Y %v\n", reply.Y)
	} else {
		fmt.Printf("call failed!\n")
	}
}

// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
func call(rpcname string, args any, reply any) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	log.Printf("calling:", err)
	return false
}
